from openai import OpenAI
from collections import defaultdict
import configparser

config = configparser.ConfigParser()
config.read('./conf.ini',encoding='utf-8')
llm_config = config['llm_config']
openai_api_key = llm_config['openai_api_key']
openai_api_base = llm_config['openai_api_base']
conf_model = llm_config['conf_model']


step3_config = config['step3']
input_cluster_path = step3_config['input_cluster_path']
output_topic_path = step3_config['output_topic_path']


def get_prompt(clusters):
    prompt = '''你的任务是：
    （1）通过给定的一系列的科技文献的聚类结果及类别下的词，概括每个类的类别。
    （2）根据所有类别的结果，概括这个领域的研究热点是什么。
    输入格式为：类别x:主题词1、主题词2、...
    输出格式为：类别x的主题标签是《主题标签》
    以下是类别及类别下的词：
    '''
    for x,y in clusters.items():
        prompt += f'类别{x}:{"、".join(y)}'
    return prompt

def read_cluster(input_cluster_path):
    clusters = defaultdict(list)
    for line in open(input_cluster_path,'r',encoding='utf-8'):
        node,cluster = line.strip().split('\t')
        clusters[cluster].append(node)
    return clusters

def myllm(prompt):
    client = OpenAI(api_key=openai_api_key, base_url=openai_api_base)

    response = client.chat.completions.create(
        model=conf_model,
        messages=[
            {"role": "system", "content": "你是一位科技情报信息的分析专家"},
            {"role": "user", "content": prompt},
        ],
        stream=False
    )
    return response.choices[0].message.content

def main():
    print("开始读取聚类数据")
    clusters = read_cluster(input_cluster_path)
    print("聚类数据读取完成")
    
    print("生成提示信息")
    prompt = get_prompt(clusters)
    print("提示信息生成完成")
    
    print("调用语言模型生成结果")
    result = myllm(prompt)
    print("结果生成完成")
    
    print("输出结果")
    print(result)
    
    print(f"将结果写入到 {output_topic_path}")
    with open(output_topic_path, 'w', encoding='utf-8') as w:
        w.write(result)
    print("结果写入完成")

if __name__ == '__main__':
    main()